"""
Command Line Interface for Easy RWKV
"""

import os
import sys
import click
from typing import Optional
from prompt_toolkit import prompt
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.history import FileHistory
from prompt_toolkit.formatted_text import FormattedText
import colorama
from colorama import Fore, Style

from .core import EasyRWKV
from .config import Config


# Initialize colorama for cross-platform colored output
colorama.init(autoreset=True, strip=False)


def supports_color() -> bool:
    """Check if terminal supports color output"""
    if hasattr(sys.stdout, 'isatty') and sys.stdout.isatty():
        # Check common environment variables
        term = os.environ.get('TERM', '').lower()
        colorterm = os.environ.get('COLORTERM', '').lower()
        
        # Check for color support indicators
        if colorterm in ('truecolor', '24bit'):
            return True
        if term in ('xterm-256color', 'xterm-color', 'linux', 'screen', 'screen-256color'):
            return True
        if 'color' in term:
            return True
        
        # On Windows, check if we're in Windows Terminal or similar
        if sys.platform == 'win32':
            wt_session = os.environ.get('WT_SESSION')
            if wt_session:
                return True
    
    return False


def print_colored(text: str, color: str = Fore.WHITE) -> None:
    """Print colored text if terminal supports it"""
    if supports_color():
        print(f"{color}{text}{Style.RESET_ALL}")
    else:
        print(text)


def print_banner():
    """Print welcome banner"""
    banner = f"""
{Fore.CYAN}╔══════════════════════════════════════════════════════════════╗
║                        Easy RWKV Chat                       ║
║              A simple wrapper for RWKV models               ║
╚══════════════════════════════════════════════════════════════╝{Style.RESET_ALL}
"""
    print(banner)


def interactive_setup() -> Config:
    """Interactive configuration setup"""
    print_colored("Welcome to Easy RWKV! Let's set up your configuration.", Fore.GREEN)
    print()
    
    config = Config()
    
    # Model path
    while True:
        model_path = input(f"{Fore.YELLOW}Enter your RWKV model path: {Style.RESET_ALL}")
        if os.path.exists(model_path):
            config.model_path = model_path
            break
        else:
            print_colored("Model file not found. Please enter a valid path.", Fore.RED)
    
    # Strategy
    strategies = ["cpu fp32", "cpu fp16", "cuda fp16", "cuda fp32"]
    print(f"\n{Fore.YELLOW}Available strategies: {', '.join(strategies)}{Style.RESET_ALL}")
    strategy = input(f"{Fore.YELLOW}Enter strategy (default: cuda fp16): {Style.RESET_ALL}") or "cuda fp16"
    if strategy in strategies:
        config.strategy = strategy
    
    # State path (optional)
    state_path = input(f"{Fore.YELLOW}Enter custom state path (optional, press Enter to skip): {Style.RESET_ALL}")
    if state_path and os.path.exists(state_path):
        config.state_path = state_path
    
    # Save config
    save_config = input(f"{Fore.YELLOW}Save this configuration? (y/n): {Style.RESET_ALL}").lower()
    if save_config in ['y', 'yes']:
        config_path = Config.get_default_config_path()
        os.makedirs(os.path.dirname(config_path), exist_ok=True)
        config.save_to_file(config_path)
        print_colored(f"Configuration saved to {config_path}", Fore.GREEN)
    
    return config


@click.group()
@click.version_option(version="0.1.0")
def cli():
    """Easy RWKV - A simple wrapper for RWKV language models"""
    pass


@cli.command()
@click.option('--model', '-m', help='''Path to RWKV model file.(BlinkDL/rwkv-7-world/RWKV-x070-World-1.5B-v3-20250127-ctx4096.pth or /path/to/your/model.pth)''')
@click.option('--config', '-c', help='Path to configuration file')
@click.option('--strategy', '-s', default='cpu fp32', help='Model strategy (default: cpu fp32)')
@click.option('--temperature', '-t', type=float, default=1.0, help='Generation temperature (default: 1.0)')
@click.option('--top-p', '-p', type=float, default=0.3, help='Top-p sampling (default: 0.3)')
@click.option('--max-tokens', type=int, default=256, help='Maximum tokens to generate (default: 256)')
@click.option('--interactive-setup', 'setup_interactive', is_flag=True, help='Run interactive setup')
def chat(model: Optional[str], config: Optional[str], strategy: str, 
         temperature: float, top_p: float, max_tokens: int, setup_interactive: bool):
    """Start interactive chat with RWKV model"""
    
    print_banner()
    
    # Load or create configuration
    if setup_interactive:
        config_obj = interactive_setup()
    elif config:
        try:
            config_obj = Config.from_file(config)
        except FileNotFoundError:
            print_colored(f"Configuration file not found: {config}", Fore.RED)
            sys.exit(1)
    else:
        # Try to load default config
        default_config_path = Config.get_default_config_path()
        if os.path.exists(default_config_path):
            config_obj = Config.from_file(default_config_path)
            print_colored(f"Loaded configuration from {default_config_path}", Fore.GREEN)
        else:
            config_obj = Config()
    
    # Override config with command line arguments
    if model:
        config_obj.model_path = model
    if strategy:
        config_obj.strategy = strategy
    if temperature:
        config_obj.temperature = temperature
    if top_p:
        config_obj.top_p = top_p
    if max_tokens:
        config_obj.max_tokens = max_tokens
    
    # Validate model path
    if not config_obj.model_path:
        print_colored("❌ 错误: 未指定模型路径", Fore.RED)
        print_colored("\n💡 解决方案:", Fore.YELLOW)
        print_colored("1. 使用 --model 参数指定模型:", Fore.WHITE)
        print_colored("   easy-rwkv chat --model 'BlinkDL/rwkv-7-world/RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.pth'", Fore.CYAN)
        print_colored("\n2. 或者查看可用模型:", Fore.WHITE)
        print_colored("   easy-rwkv models", Fore.CYAN)
        print_colored("\n3. 或者运行交互式设置:", Fore.WHITE)
        print_colored("   easy-rwkv chat --interactive-setup", Fore.CYAN)
        print_colored("\n4. 查看帮助信息:", Fore.WHITE)
        print_colored("   easy-rwkv chat --help", Fore.CYAN)
        sys.exit(1)
    
    # 如果是 HF 模型路径，不需要检查本地文件是否存在
    if not EasyRWKV._is_hf_model_name(config_obj.model_path) and not os.path.exists(config_obj.model_path):
        print_colored(f"❌ 模型文件未找到: {config_obj.model_path}", Fore.RED)
        print_colored("\n💡 解决方案:", Fore.YELLOW)
        print_colored("1. 检查文件路径是否正确", Fore.WHITE)
        print_colored("2. 或者使用HuggingFace模型名称:", Fore.WHITE)
        print_colored("   'BlinkDL/rwkv-7-world/RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.pth'", Fore.CYAN)
        print_colored("3. 查看可用模型: easy-rwkv models", Fore.WHITE)
        sys.exit(1)
    
    try:
        # Initialize RWKV
        print_colored("Initializing RWKV model...", Fore.YELLOW)
        rwkv = EasyRWKV(config_obj)
        rwkv.load_model()
        rwkv.initialize_chat()
        
        print_colored("Model loaded successfully! Type 'exit' or 'quit' to exit.", Fore.GREEN)
        print_colored("Commands: /help, /reset, /info, /config", Fore.CYAN)
        print()
        
        # Setup command completion
        commands = ['/help', '/reset', '/info', '/config', '/exit', '/quit']
        completer = WordCompleter(commands)
        
        # Setup history
        history_path = os.path.expanduser("~/.easy_rwkv/history")
        os.makedirs(os.path.dirname(history_path), exist_ok=True)
        history = FileHistory(history_path)
        
        # Main chat loop
        while True:
            try:
                # Use simple prompt - let colorama handle the coloring separately
                user_input = prompt(
                    "User: ",
                    completer=completer,
                    history=history
                ).strip()
                
                if not user_input:
                    continue
                
                # Handle commands
                if user_input.lower() in ['exit', 'quit', '/exit', '/quit']:
                    print_colored("Goodbye! 👋", Fore.GREEN)
                    break
                elif user_input == '/help':
                    print_colored("""
Available commands:
  /help     - Show this help message
  /reset    - Reset conversation state
  /info     - Show model information
  /config   - Show current configuration
  /exit     - Exit the chat
                    """, Fore.CYAN)
                    continue
                elif user_input == '/reset':
                    rwkv.reset_state()
                    rwkv.initialize_chat()
                    print_colored("Conversation state reset.", Fore.GREEN)
                    continue
                elif user_input == '/info':
                    info = rwkv.get_model_info()
                    print_colored(f"""
Model Information:
  Status: {info['status']}
  Model Path: {info['model_path']}
  Strategy: {info['strategy']}
  State Path: {info.get('state_path', 'None')}
  Tokens Processed: {info['tokens_processed']}
  Has State: {info['has_state']}
                    """, Fore.CYAN)
                    continue
                elif user_input == '/config':
                    config_dict = config_obj.to_dict()
                    print_colored("Current Configuration:", Fore.CYAN)
                    for key, value in config_dict.items():
                        print(f"  {key}: {value}")
                    continue
                
                # Generate response
                if supports_color():
                    print(f"{Fore.GREEN}Assistant: {Style.RESET_ALL}", end="", flush=True)
                else:
                    print("Assistant: ", end="", flush=True)
                
                try:
                    response = rwkv.chat(user_input)
                    print_colored(response, Fore.GREEN)
                except Exception as e:
                    print_colored(f"Error generating response: {e}", Fore.RED)
                
                print()  # Empty line for better readability
                
            except KeyboardInterrupt:
                print_colored("\nGoodbye! 👋", Fore.GREEN)
                break
            except EOFError:
                print_colored("\nGoodbye! 👋", Fore.GREEN)
                break
            except Exception as e:
                print_colored(f"An error occurred: {e}", Fore.RED)
    
    except Exception as e:
        error_msg = str(e).lower()
        
        print_colored("❌ 模型初始化失败", Fore.RED)
        print_colored("="*60, Fore.RED)
        
        # 检测CUDA相关错误
        if "cuda" in error_msg or "torch not compiled with cuda" in error_msg:
            print_colored("🚨 CUDA 配置问题:", Fore.YELLOW)
            print_colored(f"   原因: {e}", Fore.WHITE)
            print_colored("\n💡 解决方案:", Fore.GREEN)
            print_colored("1. 在 macOS 上使用 CPU 策略:", Fore.WHITE)
            print_colored("   easy-rwkv chat --model '你的模型' --strategy 'cpu fp32'", Fore.CYAN)
            print_colored("\n2. 或者设置环境变量:", Fore.WHITE)
            print_colored("   export RWKV_CUDA_ON=0", Fore.CYAN)
            print_colored("\n3. 检查当前命令是否包含 --strategy 参数", Fore.WHITE)
            
        # 检测模型文件相关错误
        elif "file not found" in error_msg or "no such file" in error_msg:
            print_colored("🚨 模型文件问题:", Fore.YELLOW)
            print_colored(f"   原因: {e}", Fore.WHITE)
            print_colored("\n💡 解决方案:", Fore.GREEN)
            print_colored("1. 检查模型路径是否正确", Fore.WHITE)
            print_colored("2. 查看可用模型:", Fore.WHITE)
            print_colored("   easy-rwkv models", Fore.CYAN)
            print_colored("3. 使用完整的HuggingFace模型路径", Fore.WHITE)
            
        # 检测权限或其他文件系统错误
        elif "permission" in error_msg or "access" in error_msg:
            print_colored("🚨 文件权限问题:", Fore.YELLOW)
            print_colored(f"   原因: {e}", Fore.WHITE)
            print_colored("\n💡 解决方案:", Fore.GREEN)
            print_colored("1. 检查文件权限", Fore.WHITE)
            print_colored("2. 尝试使用 sudo 运行", Fore.WHITE)
            print_colored("3. 检查磁盘空间", Fore.WHITE)
            
        else:
            print_colored("🚨 未知错误:", Fore.YELLOW)
            print_colored(f"   原因: {e}", Fore.WHITE)
            print_colored("\n💡 通用解决方案:", Fore.GREEN)
            print_colored("1. 确保使用正确的策略: --strategy 'cpu fp32'", Fore.WHITE)
            print_colored("2. 检查模型文件是否存在和完整", Fore.WHITE)
            print_colored("3. 查看完整帮助: easy-rwkv chat --help", Fore.WHITE)
            print_colored("4. 查看可用模型: easy-rwkv models", Fore.WHITE)
            
        print_colored("="*60, Fore.RED)
        print_colored("📋 推荐的完整命令示例:", Fore.CYAN)
        print_colored("easy-rwkv chat --model 'BlinkDL/rwkv-7-world/RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.pth' --strategy 'cpu fp32'", Fore.WHITE)
        print_colored("="*60, Fore.RED)
        
        sys.exit(1)


@cli.command()
@click.option('--model', '-m', required=True, help='''Path to RWKV model file.(BlinkDL/rwkv-7-world/RWKV-x070-World-1.5B-v3-20250127-ctx4096.pth or /path/to/your/model.pth)''')
@click.option('--message', required=True, help='Message to send to the model')
@click.option('--strategy', '-s', default='cpu fp32', help='Model strategy')
@click.option('--temperature', '-t', type=float, default=1.0, help='Generation temperature')
@click.option('--top-p', '-p', type=float, default=0.3, help='Top-p sampling')
@click.option('--max-tokens', type=int, default=256, help='Maximum tokens to generate')
def generate(model: str, message: str, strategy: str, temperature: float, 
             top_p: float, max_tokens: int):
    """Generate a single response without interactive chat"""
    
    config_obj = Config(
        model_path=model,
        strategy=strategy,
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens
    )
    
    try:
        rwkv = EasyRWKV(config_obj)
        rwkv.load_model()
        rwkv.initialize_chat()
        
        response = rwkv.chat(message)
        print(response)
        
    except Exception as e:
        error_msg = str(e).lower()
        
        print_colored("❌ 生成失败", Fore.RED)
        print_colored("="*50, Fore.RED)
        
        # 检测CUDA相关错误
        if "cuda" in error_msg or "torch not compiled with cuda" in error_msg:
            print_colored("🚨 CUDA 配置问题:", Fore.YELLOW)
            print_colored(f"   原因: {e}", Fore.WHITE)
            print_colored("\n💡 解决方案:", Fore.GREEN)
            print_colored("添加 --strategy 'cpu fp32' 参数:", Fore.WHITE)
            print_colored(f"easy-rwkv generate --model '{model}' --message '{message}' --strategy 'cpu fp32'", Fore.CYAN)
            
        # 检测模型文件相关错误
        elif "file not found" in error_msg or "no such file" in error_msg:
            print_colored("🚨 模型文件问题:", Fore.YELLOW)
            print_colored(f"   原因: {e}", Fore.WHITE)
            print_colored("\n💡 解决方案:", Fore.GREEN)
            print_colored("1. 检查模型路径是否正确", Fore.WHITE)
            print_colored("2. 查看可用模型: easy-rwkv models", Fore.WHITE)
            
        else:
            print_colored("🚨 其他错误:", Fore.YELLOW)
            print_colored(f"   原因: {e}", Fore.WHITE)
            print_colored("\n💡 建议:", Fore.GREEN)
            print_colored("使用 --strategy 'cpu fp32' 重试", Fore.WHITE)
            
        print_colored("="*50, Fore.RED)
        sys.exit(1)


@cli.command()
def setup():
    """Run interactive setup to configure Easy RWKV"""
    config = interactive_setup()
    print_colored("Setup complete! You can now run 'easy-rwkv chat' to start chatting.", Fore.GREEN)


@cli.command()
def models():
    """List all available RWKV models from BlinkDL's Hugging Face"""
    print_banner()
    print_colored("🌟 BlinkDL RWKV Models", Fore.CYAN)
    print()
    
    EasyRWKV.list_all_models()
    
    print_colored("💡 Usage examples:", Fore.YELLOW)
    print("   easy-rwkv chat --model BlinkDL/rwkv-7-world")
    print("   easy-rwkv chat --model BlinkDL/rwkv-8-pile")
    print("   easy-rwkv chat --model BlinkDL/rwkv-6-world")
    print()


@cli.command()
def clear_cache():
    """Clear the cached BlinkDL models"""
    from .core import EasyRWKV
    EasyRWKV.clear_models_cache()


def main():
    """Main entry point"""
    cli()


if __name__ == '__main__':
    main()
